"""
    SpatialIndexing

Efficient spatial data structures for GSI EnKF data assimilation system.
This module implements K-d tree based spatial indexing for fast nearest neighbor
searches, range queries, and localization radius selection. Based on the GSI/EnKF 
Fortran implementation from kdtree2.f90.

The K-d tree (k-dimensional tree) is a space-partitioning data structure that
organizes points in k-dimensional space for efficient searches. In the context
of data assimilation, it enables:

1. **Nearest Neighbor Search**: Find observations closest to analysis points
2. **Range Queries**: Find all observations within localization radius
3. **Multi-dimensional Search**: Search in (lat,lon,pressure,time) space
4. **Efficient Localization**: Support for Gaspari-Cohn and other localization functions

Mathematical Framework:
```
Search Time Complexity: O(log N) for balanced tree
Range Query: O(log N + K) where K is number of points in range
Memory Usage: O(N) where N is number of points
```

Tree Construction:
The tree is built by recursively partitioning points along alternating dimensions,
choosing split points to balance the tree for optimal search performance.
"""
module SpatialIndexing

using LinearAlgebra
using Statistics
using Random
using ..GSICoreAnalysis: AbstractAnalysisConfig

export KDTree, KDTreeNode, SpatialPoint, SearchResult
export build_kdtree, nearest_neighbors, range_search, radius_search
export multi_dimensional_search, prioritized_search
export tree_statistics, optimize_tree, destroy_tree!

"""
    SpatialPoint{T<:AbstractFloat, D}

Represents a point in D-dimensional space with associated data.
"""
struct SpatialPoint{T<:AbstractFloat, D}
    coordinates::NTuple{D, T}    # Point coordinates (lat, lon, pressure, time, etc.)
    index::Int                   # Original index in observation array
    weight::T                    # Associated weight or quality factor
    metadata::Dict{String, Any}  # Additional metadata
end

"""
    SearchResult{T<:AbstractFloat}

Result from spatial search operations.
"""
struct SearchResult{T<:AbstractFloat}
    distance::T       # Distance to query point
    point_index::Int  # Index of found point
    point_weight::T   # Weight of found point
end

"""
    KDTreeNode{T<:AbstractFloat, D}

Node in K-d tree structure.
"""
mutable struct KDTreeNode{T<:AbstractFloat, D}
    point::Union{SpatialPoint{T, D}, Nothing}  # Point stored at this node
    split_dim::Int                              # Dimension used for splitting
    split_value::T                              # Value used for splitting
    left_child::Union{KDTreeNode{T, D}, Nothing}
    right_child::Union{KDTreeNode{T, D}, Nothing}
    level::Int                                  # Depth in tree
    bounding_box::NTuple{2, NTuple{D, T}}      # (min_coords, max_coords)
end

"""
    KDTree{T<:AbstractFloat, D}

K-dimensional tree for efficient spatial searches.
"""
struct KDTree{T<:AbstractFloat, D}
    root::Union{KDTreeNode{T, D}, Nothing}
    points::Vector{SpatialPoint{T, D}}
    dimensions::Int
    size::Int
    max_depth::Int
    balance_factor::T            # Tree balance quality metric
    construction_time::T         # Time taken to build tree
    
    # Search optimization parameters
    search_epsilon::T            # Approximate search tolerance
    max_search_depth::Int        # Maximum search recursion depth
    enable_pruning::Bool         # Enable branch pruning during search
end

"""
    build_kdtree(points; method=:median, max_leaf_size=1, dimensions=nothing)

Construct K-d tree from array of spatial points.

# Arguments
- `points`: Vector of SpatialPoint objects
- `method`: Split point selection method (:median, :mean, :balanced)
- `max_leaf_size`: Maximum points per leaf node
- `dimensions`: Number of dimensions to use (default: infer from points)

# Returns
- `KDTree`: Constructed spatial index
"""
function build_kdtree(
    points::Vector{SpatialPoint{T, D}};
    method::Symbol = :median,
    max_leaf_size::Int = 1,
    dimensions::Union{Int, Nothing} = nothing
) where {T<:AbstractFloat, D}
    
    if isempty(points)
        return KDTree{T, D}(nothing, points, D, 0, 0, zero(T), zero(T), 
                           T(1e-10), 100, true)
    end
    
    start_time = time()
    effective_dims = dimensions === nothing ? D : min(dimensions, D)
    
    # Compute global bounding box
    global_bbox = compute_bounding_box(points, effective_dims)
    
    # Build tree recursively
    root = build_tree_recursive(points, effective_dims, 0, global_bbox, method, max_leaf_size)
    
    # Compute tree statistics
    max_depth = compute_tree_depth(root)
    balance_factor = compute_balance_factor(root)
    construction_time = time() - start_time
    
    return KDTree{T, D}(root, points, effective_dims, length(points), max_depth,
                        balance_factor, T(construction_time), T(1e-10), 100, true)
end

"""
    build_tree_recursive(points, dimensions, level, bbox, method, max_leaf_size)

Recursively construct K-d tree nodes.
"""
function build_tree_recursive(
    points::Vector{SpatialPoint{T, D}},
    dimensions::Int,
    level::Int,
    bbox::NTuple{2, NTuple{D, T}},
    method::Symbol,
    max_leaf_size::Int
) where {T<:AbstractFloat, D}
    
    if length(points) <= max_leaf_size
        # Create leaf node
        point = length(points) == 1 ? points[1] : nothing
        return KDTreeNode{T, D}(point, -1, zero(T), nothing, nothing, level, bbox)
    end
    
    # Choose splitting dimension (cycle through dimensions)
    split_dim = (level % dimensions) + 1
    
    # Select split value based on method
    split_value = select_split_value(points, split_dim, method)
    
    # Partition points
    left_points = SpatialPoint{T, D}[]
    right_points = SpatialPoint{T, D}[]
    
    for point in points
        if point.coordinates[split_dim] <= split_value
            push!(left_points, point)
        else
            push!(right_points, point)
        end
    end
    
    # Handle edge case where all points go to one side
    if isempty(left_points) || isempty(right_points)
        # Split points evenly
        sorted_indices = sortperm(points, by = p -> p.coordinates[split_dim])
        mid_point = length(points) ÷ 2
        left_points = points[sorted_indices[1:mid_point]]
        right_points = points[sorted_indices[mid_point+1:end]]
        split_value = (left_points[end].coordinates[split_dim] + 
                      right_points[1].coordinates[split_dim]) / 2
    end
    
    # Compute child bounding boxes
    left_bbox, right_bbox = split_bounding_box(bbox, split_dim, split_value)
    
    # Recursively build child nodes
    left_child = build_tree_recursive(left_points, dimensions, level + 1, 
                                    left_bbox, method, max_leaf_size)
    right_child = build_tree_recursive(right_points, dimensions, level + 1,
                                     right_bbox, method, max_leaf_size)
    
    return KDTreeNode{T, D}(nothing, split_dim, split_value, left_child, 
                           right_child, level, bbox)
end

"""
    select_split_value(points, dimension, method)

Select split value for tree construction.
"""
function select_split_value(
    points::Vector{SpatialPoint{T, D}},
    dimension::Int,
    method::Symbol
) where {T<:AbstractFloat, D}
    
    values = [p.coordinates[dimension] for p in points]
    
    if method == :median
        return Statistics.median(values)
    elseif method == :mean
        return Statistics.mean(values)
    elseif method == :balanced
        # Choose value that balances tree better
        sorted_values = sort(values)
        mid_idx = length(sorted_values) ÷ 2
        return sorted_values[mid_idx]
    else
        error("Unknown split method: $method")
    end
end

"""
    nearest_neighbors(tree, query_point, k; distance_metric=:euclidean)

Find k nearest neighbors to query point.

# Arguments
- `tree`: KDTree to search
- `query_point`: Query coordinates as NTuple or Vector
- `k`: Number of neighbors to find
- `distance_metric`: Distance metric (:euclidean, :manhattan, :haversine)

# Returns
- `Vector{SearchResult}`: k nearest neighbors sorted by distance
"""
function nearest_neighbors(
    tree::KDTree{T, D},
    query_point::Union{NTuple{D, T}, Vector{T}},
    k::Int;
    distance_metric::Symbol = :euclidean
) where {T<:AbstractFloat, D}
    
    if tree.root === nothing
        return SearchResult{T}[]
    end
    
    query_coords = query_point isa Vector ? tuple(query_point...) : query_point
    
    # Priority queue to maintain k best candidates
    candidates = SearchResult{T}[]
    
    # Recursive search
    knn_search_recursive!(candidates, tree.root, query_coords, k, distance_metric, tree)
    
    # Sort by distance and return top k
    sort!(candidates, by = r -> r.distance)
    return candidates[1:min(k, length(candidates))]
end

"""
    knn_search_recursive!(candidates, node, query_coords, k, distance_metric, tree)

Recursive k-nearest neighbors search implementation.
"""
function knn_search_recursive!(
    candidates::Vector{SearchResult{T}},
    node::KDTreeNode{T, D},
    query_coords::NTuple{D, T},
    k::Int,
    distance_metric::Symbol,
    tree::KDTree{T, D}
) where {T<:AbstractFloat, D}
    
    if node.point !== nothing
        # Leaf node - evaluate point
        distance = compute_distance(query_coords, node.point.coordinates, distance_metric)
        result = SearchResult(distance, node.point.index, node.point.weight)
        
        add_candidate!(candidates, result, k)
        return
    end
    
    # Internal node - decide which subtree to search first
    split_dim = node.split_dim
    query_value = query_coords[split_dim]
    
    # Determine primary and secondary subtrees
    if query_value <= node.split_value
        primary_child = node.left_child
        secondary_child = node.right_child
    else
        primary_child = node.right_child
        secondary_child = node.left_child
    end
    
    # Search primary subtree
    if primary_child !== nothing
        knn_search_recursive!(candidates, primary_child, query_coords, k, 
                            distance_metric, tree)
    end
    
    # Check if we need to search secondary subtree
    if secondary_child !== nothing && should_search_subtree(candidates, node, 
                                                          query_coords, k, distance_metric)
        knn_search_recursive!(candidates, secondary_child, query_coords, k,
                            distance_metric, tree)
    end
end

"""
    range_search(tree, query_point, radius; distance_metric=:euclidean)

Find all points within specified radius of query point.

# Arguments
- `tree`: KDTree to search
- `query_point`: Query coordinates
- `radius`: Search radius
- `distance_metric`: Distance metric to use

# Returns
- `Vector{SearchResult}`: All points within radius
"""
function range_search(
    tree::KDTree{T, D},
    query_point::Union{NTuple{D, T}, Vector{T}},
    radius::T;
    distance_metric::Symbol = :euclidean
) where {T<:AbstractFloat, D}
    
    if tree.root === nothing
        return SearchResult{T}[]
    end
    
    query_coords = query_point isa Vector ? tuple(query_point...) : query_point
    results = SearchResult{T}[]
    
    range_search_recursive!(results, tree.root, query_coords, radius, distance_metric)
    
    return results
end

"""
    range_search_recursive!(results, node, query_coords, radius, distance_metric)

Recursive range search implementation.
"""
function range_search_recursive!(
    results::Vector{SearchResult{T}},
    node::KDTreeNode{T, D},
    query_coords::NTuple{D, T},
    radius::T,
    distance_metric::Symbol
) where {T<:AbstractFloat, D}
    
    if node.point !== nothing
        # Leaf node - check distance
        distance = compute_distance(query_coords, node.point.coordinates, distance_metric)
        if distance <= radius
            push!(results, SearchResult(distance, node.point.index, node.point.weight))
        end
        return
    end
    
    # Internal node - check which subtrees to search
    split_dim = node.split_dim
    split_value = node.split_value
    query_value = query_coords[split_dim]
    
    # Distance to splitting hyperplane
    plane_distance = abs(query_value - split_value)
    
    # Search appropriate subtrees
    if query_value <= split_value
        # Query point is on left side
        if node.left_child !== nothing
            range_search_recursive!(results, node.left_child, query_coords, 
                                  radius, distance_metric)
        end
        
        # Check if sphere intersects right side
        if plane_distance <= radius && node.right_child !== nothing
            range_search_recursive!(results, node.right_child, query_coords,
                                  radius, distance_metric)
        end
    else
        # Query point is on right side
        if node.right_child !== nothing
            range_search_recursive!(results, node.right_child, query_coords,
                                  radius, distance_metric)
        end
        
        # Check if sphere intersects left side
        if plane_distance <= radius && node.left_child !== nothing
            range_search_recursive!(results, node.left_child, query_coords,
                                  radius, distance_metric)
        end
    end
end

"""
    radius_search(tree, query_point, min_radius, max_radius)

Find all points within annular region between min_radius and max_radius.
"""
function radius_search(
    tree::KDTree{T, D},
    query_point::Union{NTuple{D, T}, Vector{T}},
    min_radius::T,
    max_radius::T;
    distance_metric::Symbol = :euclidean
) where {T<:AbstractFloat, D}
    
    # Get all points within max_radius
    all_results = range_search(tree, query_point, max_radius; distance_metric = distance_metric)
    
    # Filter to only include points beyond min_radius
    return filter(r -> r.distance >= min_radius, all_results)
end

"""
    multi_dimensional_search(tree, query_ranges)

Search for points within multi-dimensional rectangular region.

# Arguments
- `tree`: KDTree to search
- `query_ranges`: Vector of (min, max) tuples for each dimension

# Returns
- `Vector{SearchResult}`: All points within specified ranges
"""
function multi_dimensional_search(
    tree::KDTree{T, D},
    query_ranges::Vector{Tuple{T, T}}
) where {T<:AbstractFloat, D}
    
    if tree.root === nothing || length(query_ranges) > D
        return SearchResult{T}[]
    end
    
    results = SearchResult{T}[]
    multi_dim_search_recursive!(results, tree.root, query_ranges)
    
    return results
end

"""
    multi_dim_search_recursive!(results, node, query_ranges)

Recursive multi-dimensional range search.
"""
function multi_dim_search_recursive!(
    results::Vector{SearchResult{T}},
    node::KDTreeNode{T, D},
    query_ranges::Vector{Tuple{T, T}}
) where {T<:AbstractFloat, D}
    
    if node.point !== nothing
        # Leaf node - check if point is within all ranges
        within_range = true
        for (dim, (min_val, max_val)) in enumerate(query_ranges)
            if dim > D
                break
            end
            coord = node.point.coordinates[dim]
            if coord < min_val || coord > max_val
                within_range = false
                break
            end
        end
        
        if within_range
            # Distance not meaningful for range queries, use 0
            push!(results, SearchResult(zero(T), node.point.index, node.point.weight))
        end
        return
    end
    
    # Internal node - check intersection with query ranges
    split_dim = node.split_dim
    split_value = node.split_value
    
    if split_dim <= length(query_ranges)
        min_range, max_range = query_ranges[split_dim]
        
        # Check left subtree
        if split_value >= min_range && node.left_child !== nothing
            multi_dim_search_recursive!(results, node.left_child, query_ranges)
        end
        
        # Check right subtree
        if split_value <= max_range && node.right_child !== nothing
            multi_dim_search_recursive!(results, node.right_child, query_ranges)
        end
    else
        # Split dimension not constrained - search both subtrees
        if node.left_child !== nothing
            multi_dim_search_recursive!(results, node.left_child, query_ranges)
        end
        if node.right_child !== nothing
            multi_dim_search_recursive!(results, node.right_child, query_ranges)
        end
    end
end

"""
    compute_distance(point1, point2, metric)

Compute distance between two points using specified metric.
"""
function compute_distance(
    point1::NTuple{D, T},
    point2::NTuple{D, T},
    metric::Symbol
) where {T<:AbstractFloat, D}
    
    if metric == :euclidean
        sum_squares = zero(T)
        for i in 1:D
            diff = point1[i] - point2[i]
            sum_squares += diff * diff
        end
        return sqrt(sum_squares)
        
    elseif metric == :manhattan
        sum_abs = zero(T)
        for i in 1:D
            sum_abs += abs(point1[i] - point2[i])
        end
        return sum_abs
        
    elseif metric == :haversine
        # Great circle distance for lat/lon coordinates
        if D < 2
            error("Haversine distance requires at least 2 dimensions (lat, lon)")
        end
        
        lat1, lon1 = deg2rad(point1[1]), deg2rad(point1[2])
        lat2, lon2 = deg2rad(point2[1]), deg2rad(point2[2])
        
        dlat = lat2 - lat1
        dlon = lon2 - lon1
        
        a = sin(dlat/2)^2 + cos(lat1) * cos(lat2) * sin(dlon/2)^2
        c = 2 * atan(sqrt(a), sqrt(1-a))
        
        # Earth radius in km
        R = T(6371.0)
        return R * c
        
    else
        error("Unknown distance metric: $metric")
    end
end

"""
    add_candidate!(candidates, new_result, k)

Add new candidate to k-NN results, maintaining size constraint.
"""
function add_candidate!(candidates::Vector{SearchResult{T}}, new_result::SearchResult{T}, k::Int) where T
    if length(candidates) < k
        push!(candidates, new_result)
        sort!(candidates, by = r -> r.distance)
    elseif new_result.distance < candidates[end].distance
        candidates[end] = new_result
        sort!(candidates, by = r -> r.distance)
    end
end

"""
    should_search_subtree(candidates, node, query_coords, k, distance_metric)

Determine if subtree should be searched based on current best candidates.
"""
function should_search_subtree(
    candidates::Vector{SearchResult{T}},
    node::KDTreeNode{T, D},
    query_coords::NTuple{D, T},
    k::Int,
    distance_metric::Symbol
) where {T<:AbstractFloat, D}
    
    if length(candidates) < k
        return true
    end
    
    # Distance to splitting hyperplane
    split_dim = node.split_dim
    plane_distance = abs(query_coords[split_dim] - node.split_value)
    
    # Compare with worst current candidate
    worst_distance = candidates[end].distance
    
    return plane_distance <= worst_distance
end

"""
    compute_bounding_box(points, dimensions)

Compute axis-aligned bounding box for set of points.
"""
function compute_bounding_box(
    points::Vector{SpatialPoint{T, D}},
    dimensions::Int
) where {T<:AbstractFloat, D}
    
    if isempty(points)
        return (ntuple(i -> zero(T), D), ntuple(i -> zero(T), D))
    end
    
    min_coords = collect(points[1].coordinates)
    max_coords = collect(points[1].coordinates)
    
    for point in points[2:end]
        for dim in 1:min(dimensions, D)
            coord = point.coordinates[dim]
            if coord < min_coords[dim]
                min_coords[dim] = coord
            elseif coord > max_coords[dim]
                max_coords[dim] = coord
            end
        end
    end
    
    return (tuple(min_coords[1:D]...), tuple(max_coords[1:D]...))
end

"""
    split_bounding_box(bbox, split_dim, split_value)

Split bounding box along specified dimension.
"""
function split_bounding_box(
    bbox::NTuple{2, NTuple{D, T}},
    split_dim::Int,
    split_value::T
) where {T<:AbstractFloat, D}
    
    min_coords, max_coords = bbox
    
    # Left child bounding box
    left_max = collect(max_coords)
    left_max[split_dim] = split_value
    left_bbox = (min_coords, tuple(left_max...))
    
    # Right child bounding box
    right_min = collect(min_coords)
    right_min[split_dim] = split_value
    right_bbox = (tuple(right_min...), max_coords)
    
    return left_bbox, right_bbox
end

"""
    compute_tree_depth(node)

Compute maximum depth of tree.
"""
function compute_tree_depth(node::Union{KDTreeNode{T, D}, Nothing}) where {T, D}
    if node === nothing
        return 0
    end
    
    left_depth = compute_tree_depth(node.left_child)
    right_depth = compute_tree_depth(node.right_child)
    
    return 1 + max(left_depth, right_depth)
end

"""
    compute_balance_factor(node)

Compute tree balance quality metric (0 = perfectly balanced, 1 = completely unbalanced).
"""
function compute_balance_factor(node::Union{KDTreeNode{T, D}, Nothing}) where {T, D}
    if node === nothing
        return zero(T)
    end
    
    left_depth = compute_tree_depth(node.left_child)
    right_depth = compute_tree_depth(node.right_child)
    
    max_depth = max(left_depth, right_depth)
    if max_depth == 0
        return zero(T)
    end
    
    return T(abs(left_depth - right_depth)) / T(max_depth)
end

"""
    tree_statistics(tree)

Compute comprehensive tree statistics.
"""
function tree_statistics(tree::KDTree{T, D}) where {T, D}
    return Dict(
        "size" => tree.size,
        "dimensions" => tree.dimensions,
        "max_depth" => tree.max_depth,
        "balance_factor" => tree.balance_factor,
        "construction_time" => tree.construction_time,
        "memory_usage_MB" => sizeof(tree) / (1024^2)
    )
end

"""
    optimize_tree(tree; rebalance_threshold=0.3)

Optimize tree structure if balance factor exceeds threshold.
"""
function optimize_tree(
    tree::KDTree{T, D};
    rebalance_threshold::T = T(0.3)
) where {T<:AbstractFloat, D}
    
    if tree.balance_factor > rebalance_threshold
        # Rebuild tree with better balance
        return build_kdtree(tree.points; method = :balanced, max_leaf_size = 1)
    end
    
    return tree
end

"""
    destroy_tree!(tree)

Clean up tree memory (placeholder for explicit memory management if needed).
"""
function destroy_tree!(tree::KDTree{T, D}) where {T, D}
    # Julia handles garbage collection automatically
    # This function exists for API compatibility with C/Fortran versions
    return nothing
end

"""
    prioritized_search(tree, query_point, priorities, k)

Search with priority weighting (e.g., observation quality, temporal proximity).
"""
function prioritized_search(
    tree::KDTree{T, D},
    query_point::Union{NTuple{D, T}, Vector{T}},
    priorities::Dict{Int, T},
    k::Int;
    distance_metric::Symbol = :euclidean
) where {T<:AbstractFloat, D}
    
    # Find larger set of candidates
    candidates = nearest_neighbors(tree, query_point, min(k * 3, tree.size); 
                                 distance_metric = distance_metric)
    
    # Reweight by priorities
    for candidate in candidates
        priority = get(priorities, candidate.point_index, one(T))
        # Combine distance and priority (lower is better)
        candidate = SearchResult(candidate.distance / priority, 
                               candidate.point_index, candidate.point_weight)
    end
    
    # Sort by combined score and return top k
    sort!(candidates, by = r -> r.distance)
    return candidates[1:min(k, length(candidates))]
end

end # module SpatialIndexing